/*
 * Copyright (c) 2013-2018 上海汇数数据. All rights reserved.
 * @(#) KeyKt.java 2018-08-06 17:26
 */

package cn.jh.common.core.crypto;

import cn.jh.common.core.base.Base64Kt;
import com.google.common.base.Splitter;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.openssl.PEMParser;
import org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter;
import org.bouncycastle.openssl.jcajce.JcaPEMWriter;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.StringReader;
import java.io.StringWriter;
import java.io.UncheckedIOException;
import java.security.*;
import java.security.spec.AlgorithmParameterSpec;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;

import static com.google.common.base.Preconditions.checkArgument;


/**
 * 密钥工厂帮助工具类。
 *
 * @author Fuchun
 * @since 1.0
 */
public class KeyKt {

    private static final Logger LOGGER = LoggerFactory.getLogger(KeyKt.class);

    private static final Pattern NEXT_LINE = Pattern.compile("[\r\n]");
    private static final Splitter NL_SPLITTER = Splitter.on(NEXT_LINE).omitEmptyStrings();

    public static String formatPEM(Key key) {
        StringWriter writer;
        try (JcaPEMWriter pemWriter = new JcaPEMWriter((writer = new StringWriter()))) {

            pemWriter.writeObject(key);
            pemWriter.flush();

            return writer.toString();
        } catch (IOException ex) {
            throw new UncheckedIOException(ex);
        }

    }

    public static PrivateKey toPrivateKeyFromPEM(String pemKey, @NotNull String algorithm) throws CryptoException {
        byte[] keyBytes;
        if (!pemKey.contains("\r\n") && !pemKey.contains("\n")) {
            keyBytes = Base64Kt.decode(pemKey);
            checkArgument(keyBytes != null && keyBytes.length > 0,
                    "The pemKey cannot decode to use Base64 format.");
            return toPrivateKey(keyBytes, algorithm);
        }
        Iterable<String> lines = NL_SPLITTER.split(pemKey);

        String b64Key = StreamSupport.stream(lines.spliterator(), false)
                .filter(line -> !line.startsWith("--"))
                .collect(Collectors.joining(""));
        keyBytes = Base64Kt.decode(b64Key);
        checkArgument(keyBytes != null && keyBytes.length > 0,
                "The pemKey cannot decode to use Base64 format.");
        return toPrivateKey(keyBytes, algorithm);
    }

    public static PublicKey toPublicKeyFromPEM(String pemKey) throws CryptoException {

        try (PEMParser parser = new PEMParser(new StringReader(pemKey))) {
            SubjectPublicKeyInfo publicKeyInfo = (SubjectPublicKeyInfo) parser.readObject();
            return new JcaPEMKeyConverter().getPublicKey(publicKeyInfo);
        } catch (IOException ex) {
            throw new UncheckedIOException(ex);
        }
    }

    /**
     * 将指定的密钥数据转换为指定非对称加密算法的私钥。
     *
     * @param key       非对称加密算法密钥。
     * @param algorithm 非对称加密算法名称，例如：RSA，DSA
     * @return 返回指定算法下的给定数据表示的非对称加密算法的私钥。
     * @throws CryptoException 如果指定算法下给定的密钥错误（包括但不限于密钥长度，格式，编码等）。
     */
    public static PrivateKey toPrivateKey(@NotNull byte[] key, @NotNull String algorithm) throws CryptoException {
        checkArgument(key.length > 0, "The key bytes must be not empty.");
        checkArgument(algorithm.length() > 0, "The algorithm must be not empty.");

        KeyFactory kf = getBcKeyFactory(algorithm);
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(key);
        try {
            return kf.generatePrivate(keySpec);
        } catch (InvalidKeySpecException ex) {
            LOGGER.error("Invalid private key spec: {}", ex.getLocalizedMessage() == null ?
                    ex.getMessage() : ex.getLocalizedMessage());
            throw new CryptoException("The private key is invalid : \n" + Base64Kt.encodeToMime(key));
        }
    }

    /**
     * 将指定的密钥数据转换为指定非对称加密算法的公钥。
     *
     * @param key       非对称加密算法密钥。
     * @param algorithm 非对称加密算法名称，例如：RSA，DSA
     * @return 返回指定算法下的给定数据表示的非对称加密算法的公钥。
     * @throws CryptoException 如果指定算法下给定的密钥错误（包括但不限于密钥长度，格式，编码等）。
     */
    public static PublicKey toPublicKey(@NotNull byte[] key, @NotNull String algorithm) throws CryptoException {
        checkArgument(key.length > 0, "The key bytes must be not empty.");
        checkArgument(algorithm.length() > 0, "The algorithm must be not empty.");
        KeyFactory kf = getBcKeyFactory(algorithm);
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(key);
        try {
            return kf.generatePublic(keySpec);
        } catch (InvalidKeySpecException ex) {
            LOGGER.error("Invalid public key: {}", ex.getLocalizedMessage() == null ?
                    ex.getMessage() : ex.getLocalizedMessage());
            throw new CryptoException("The public key is invalid: \n" + Base64Kt.encodeToMime(key));
        }
    }

    protected static KeyFactory getKeyFactory(String algorithm, @Nullable String provider) {
        try {
            if (provider == null) {
                return KeyFactory.getInstance(algorithm);
            } else {
                return KeyFactory.getInstance(algorithm, provider);
            }
        } catch (NoSuchAlgorithmException | NoSuchProviderException ex) {
            throw new IllegalArgumentException(ex);
        }
    }

    public static KeyFactory getBcKeyFactory(String algorithm) {
        return getKeyFactory(algorithm, CipherKt.BOUNCY_CASTLE);
    }

    public static KeyFactory getJceKeyFactory(String algorithm) {
        return getKeyFactory(algorithm, null);
    }

    protected static KeyPairGenerator getKpGen(String algorithm, @Nullable String provider) {
        try {
            KeyPairGenerator kpGen;
            if (provider == null) {
                kpGen = KeyPairGenerator.getInstance(algorithm);
            } else {
                kpGen = KeyPairGenerator.getInstance(algorithm, provider);
            }
            return kpGen;
        } catch (NoSuchAlgorithmException | NoSuchProviderException ex) {
            throw new IllegalArgumentException(ex);
        }
    }

    public static KeyPairGenerator getAndInitKpGen(String algorithm, @Nullable String provider, KeyBit keyBit) {
        KeyPairGenerator kpGen = getKpGen(algorithm, provider);
        kpGen.initialize(keyBit.getBit(), new SecureRandom());
        return kpGen;
    }

    public static KeyPairGenerator getAndInitKpGen(
            String algorithm, @Nullable String provider, AlgorithmParameterSpec parameterSpec) {
        KeyPairGenerator kpGen = getKpGen(algorithm, provider);
        try {
            kpGen.initialize(parameterSpec, new SecureRandom());
        } catch (InvalidAlgorithmParameterException ex) {
            throw new IllegalArgumentException(ex);
        }
        return kpGen;
    }

    public static KeyPairGenerator getAndInitJceKpGen(String algorithm, KeyBit keyBit) {
        return getAndInitKpGen(algorithm, null, keyBit);
    }

    public static KeyPairGenerator getAndInitJceKpGen(String algorithm, AlgorithmParameterSpec parameterSpec) {
        return getAndInitKpGen(algorithm, null, parameterSpec);
    }

    public static KeyPairGenerator getAndInitBcKpGen(String algorithm, KeyBit keyBit) {
        return getAndInitKpGen(algorithm, CipherKt.BOUNCY_CASTLE, keyBit);
    }

    public static KeyPairGenerator getAndInitBcKpGen(String algorithm, AlgorithmParameterSpec parameterSpec) {
        return getAndInitKpGen(algorithm, CipherKt.BOUNCY_CASTLE, parameterSpec);
    }

    private KeyKt() {
    }
}
